import cv2
import torch
import numpy as np
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.general import non_max_suppression, scale_boxes


def load_yolov5_model(weights='yolov5s.pt'):
    """加载YOLOv5模型"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = DetectMultiBackend(weights, device=device)
    return model


def detect_people(frame, model, conf_threshold=0.5):
    """使用YOLOv5模型检测图像中的人"""
    frame = torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0
    if frame.ndimension() == 3:
        frame = frame.unsqueeze(0)
    pred = model(frame)[0]
    pred = non_max_suppression(pred, conf_threshold, 0.4)
    boxes = scale_boxes(pred[0][:, :4], frame.shape[1:], 1.0)[0]
    return boxes


def capture_and_detect():
    """开启摄像头，检测到人时进行截图"""
    # 加载YOLOv5模型
    model = load_yolov5_model()

    # 打开摄像头
    cap = cv2.VideoCapture(0)  # 0代表默认摄像头

    while True:
        # 读取一帧
        ret, frame = cap.read()
        if not ret:
            break

        # 检测人物
        detections = detect_people(frame, model)

        # 如果检测到人，保存截图
        if len(detections) > 0:
            cv2.imwrite("detected_person_screenshot.jpg", frame)
            print("检测到人，已截图保存。")

        # 显示帧（可选）
        # cv2.imshow("Camera", frame)
        # if cv2.waitKey(1) & 0xFF == ord('q'):  # 按'q'退出
        #     break

    # 释放资源
    cap.release()
    # cv2.destroyAllWindows()  # 如果上面有显示窗口，记得释放


if __name__ == "__main__":
    capture_and_detect()